"""
百度OCR工具模块
提供图片和PDF文字识别功能，集成到LangChain工具链
"""

import base64
import logging
from typing import Optional, Dict, Any
from pathlib import Path

logger = logging.getLogger(__name__)

try:
    from aip import AipOcr
except ImportError as e:
    AipOcr = None
    import sys
    logger.error(f"导入AipOcr失败: {e}")
    logger.error(f"Python路径: {sys.executable}")
    logger.error(f"sys.path: {sys.path}")

from langchain.tools import BaseTool
from pydantic import Field

# 导入OCR使用次数跟踪
try:
    from .ocr_usage_tracker import increment_ocr_usage, check_usage_limit
except ImportError:
    # 如果导入失败，提供空实现
    def increment_ocr_usage():
        return 0
    def check_usage_limit():
        return True, ""


class BaiduOCRClient:
    """百度OCR客户端封装类"""
    
    def __init__(self, app_id: str, api_key: str, secret_key: str):
        """
        初始化百度OCR客户端
        
        Args:
            app_id: 百度云应用ID
            api_key: 百度云API Key
            secret_key: 百度云Secret Key
        """
        if AipOcr is None:
            import sys
            error_msg = (
                f"未安装百度AI SDK，请运行: pip install baidu-aip chardet\n"
                f"当前Python: {sys.executable}\n"
                f"请确保在正确的Python环境中安装"
            )
            raise ImportError(error_msg)
        
        self.client = AipOcr(app_id, api_key, secret_key)
        logger.info("百度OCR客户端初始化成功")
    
    def _read_file_as_base64(self, file_path: str) -> bytes:
        """
        读取文件并转换为字节流
        
        Args:
            file_path: 文件路径
            
        Returns:
            文件字节内容
        """
        with open(file_path, 'rb') as f:
            return f.read()
    
    def recognize_general(
        self,
        file_path: str,
        options: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        通用文字识别（标准版）
        
        Args:
            file_path: 图片文件路径
            options: 可选参数，如 language_type, detect_direction 等
            
        Returns:
            识别的文字内容
        """
        try:
            image = self._read_file_as_base64(file_path)
            
            if options is None:
                options = {}
            
            # 调用百度OCR API
            result = self.client.basicGeneral(image, options)
            
            # 提取文字内容
            if 'words_result' in result:
                text_lines = [item['words'] for item in result['words_result']]
                recognized_text = '\n'.join(text_lines)
                
                logger.info(f"OCR识别成功，识别出 {len(text_lines)} 行文字")
                return recognized_text
            else:
                error_msg = result.get('error_msg', '未知错误')
                logger.error(f"OCR识别失败: {error_msg}")
                return f"OCR识别失败: {error_msg}"
                
        except Exception as e:
            logger.error(f"OCR识别异常: {e}")
            return f"OCR识别异常: {str(e)}"
    
    def recognize_accurate(
        self,
        file_path: str,
        options: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        通用文字识别（高精度版）
        
        Args:
            file_path: 图片文件路径
            options: 可选参数
            
        Returns:
            识别的文字内容
        """
        try:
            image = self._read_file_as_base64(file_path)
            
            if options is None:
                options = {}
            
            # 调用高精度OCR API
            result = self.client.basicAccurate(image, options)
            
            # 提取文字内容
            if 'words_result' in result:
                text_lines = [item['words'] for item in result['words_result']]
                recognized_text = '\n'.join(text_lines)
                
                logger.info(f"高精度OCR识别成功，识别出 {len(text_lines)} 行文字")
                return recognized_text
            else:
                error_msg = result.get('error_msg', '未知错误')
                logger.error(f"高精度OCR识别失败: {error_msg}")
                return f"OCR识别失败: {error_msg}"
                
        except Exception as e:
            logger.error(f"高精度OCR识别异常: {e}")
            return f"OCR识别异常: {str(e)}"
    
    def recognize_pdf(
        self,
        file_path: str,
        options: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        PDF文件识别
        
        Args:
            file_path: PDF文件路径
            options: 可选参数
            
        Returns:
            识别的文字内容
        """
        try:
            pdf_file = self._read_file_as_base64(file_path)
            
            if options is None:
                options = {}
            
            # 调用PDF识别API
            result = self.client.pdf(pdf_file, options)
            
            # 提取文字内容
            if 'words_result' in result:
                text_lines = [item['words'] for item in result['words_result']]
                recognized_text = '\n'.join(text_lines)
                
                logger.info(f"PDF识别成功，识别出 {len(text_lines)} 行文字")
                return recognized_text
            else:
                error_msg = result.get('error_msg', '未知错误')
                logger.error(f"PDF识别失败: {error_msg}")
                return f"PDF识别失败: {error_msg}"
                
        except Exception as e:
            logger.error(f"PDF识别异常: {e}")
            return f"PDF识别异常: {str(e)}"


class BaiduOCRTool(BaseTool):
    """
    百度OCR工具 - LangChain Tool封装
    可被LangChain Agent调用
    """
    
    name: str = "baidu_ocr"
    description: str = (
        "用于识别图片或PDF文件中的文字内容。"
        "输入应该是文件路径。"
        "支持PNG、JPG、JPEG、PDF等格式。"
    )
    
    # 使用Field标记为私有属性
    ocr_client: Any = Field(default=None, exclude=True)
    ocr_type: str = Field(default="accurate_basic", exclude=True)
    
    def __init__(self, app_id: str, api_key: str, secret_key: str, ocr_type: str = "accurate_basic"):
        """
        初始化OCR工具
        
        Args:
            app_id: 百度云应用ID
            api_key: 百度云API Key
            secret_key: 百度云Secret Key
            ocr_type: OCR类型 (general_basic, accurate_basic)
        """
        super().__init__()
        self.ocr_client = BaiduOCRClient(app_id, api_key, secret_key)
        self.ocr_type = ocr_type
        logger.info(f"BaiduOCRTool初始化成功，OCR类型: {ocr_type}")
    
    def _run(self, file_path: str) -> str:
        """
        执行OCR识别
        
        Args:
            file_path: 文件路径
            
        Returns:
            识别的文字内容
        """
        logger.info(f"开始OCR识别: {file_path}")
        
        # 检查文件是否存在
        if not Path(file_path).exists():
            error_msg = f"文件不存在: {file_path}"
            logger.error(error_msg)
            return error_msg
        
        # 检查OCR配额
        can_use, warning_msg = check_usage_limit()
        if not can_use:
            logger.warning(warning_msg)
            return f"错误: {warning_msg}"
        
        # 根据文件类型选择识别方法
        file_ext = Path(file_path).suffix.lower()
        
        try:
            if file_ext == '.pdf':
                result = self.ocr_client.recognize_pdf(file_path)
            elif self.ocr_type == "accurate_basic":
                result = self.ocr_client.recognize_accurate(file_path)
            else:
                result = self.ocr_client.recognize_general(file_path)
            
            # 记录OCR使用次数
            usage_count = increment_ocr_usage()
            
            logger.info(f"OCR识别完成，内容长度: {len(result)} 字符，今日使用: {usage_count}/200")
            return result
            
        except Exception as e:
            error_msg = f"OCR识别失败: {str(e)}"
            logger.error(error_msg)
            return error_msg
    
    async def _arun(self, file_path: str) -> str:
        """异步执行（暂时调用同步方法）"""
        return self._run(file_path)


def create_baidu_ocr_tool(
    app_id: str,
    api_key: str,
    secret_key: str,
    ocr_type: str = "accurate_basic"
) -> BaiduOCRTool:
    """
    创建百度OCR工具实例的工厂函数
    
    Args:
        app_id: 百度云应用ID
        api_key: 百度云API Key
        secret_key: 百度云Secret Key
        ocr_type: OCR类型
        
    Returns:
        BaiduOCRTool实例
    """
    return BaiduOCRTool(
        app_id=app_id,
        api_key=api_key,
        secret_key=secret_key,
        ocr_type=ocr_type
    )

